{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tool and Function Calling\n",
    "\n",
    "This guide demonstrates how to use RTP-LLM’s [Function calling](https://platform.openai.com/docs/guides/function-calling) functionality."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OpenAI Compatible API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently, RTP-LLM supports Function Calling capabilities for the following models:\n",
    "\n",
    "- qwen25/qwen3: Qwen 2.5 (e.g. Qwen/Qwen2.5-1.5B-Instruct, Qwen/Qwen2.5-7B-Instruct) and QwQ (i.e. Qwen/QwQ-32B), Qwen3-235B-A22B, etc.\n",
    "- Kimi_k2\n",
    "- Qwen3-Coder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Tools for Function Call\n",
    "Below is a Python snippet that shows how to define a tool as a dictionary. The dictionary includes a tool name, a description, and property defined Parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define tools\n",
    "tools = [\n",
    "    {\n",
    "        \"type\": \"function\",\n",
    "        \"function\": {\n",
    "            \"name\": \"get_current_weather\",\n",
    "            \"description\": \"Get the current weather in a given location\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"city\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The city to find the weather for, e.g. 'San Francisco'\",\n",
    "                    },\n",
    "                    \"state\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"the two-letter abbreviation for the state that the city is\"\n",
    "                        \" in, e.g. 'CA' which would mean 'California'\",\n",
    "                    },\n",
    "                    \"unit\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The unit to fetch the temperature in\",\n",
    "                        \"enum\": [\"celsius\", \"fahrenheit\"],\n",
    "                    },\n",
    "                },\n",
    "                \"required\": [\"city\", \"state\", \"unit\"],\n",
    "            },\n",
    "        },\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_messages():\n",
    "    return [\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": \"What's the weather like in Boston today? Output a reasoning before act, then use the tools to help you.\",\n",
    "        }\n",
    "    ]\n",
    "\n",
    "\n",
    "messages = get_messages()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize the Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize OpenAI-like client\n",
    "client = OpenAI(api_key=\"None\", base_url=f\"http://0.0.0.0:{port}/v1\")\n",
    "model_name = client.models.list().data[0].id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  Non-Streaming Request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Non-streaming mode test\n",
    "response_non_stream = client.chat.completions.create(\n",
    "    model=model_name,\n",
    "    messages=messages,\n",
    "    temperature=0,\n",
    "    top_p=0.95,\n",
    "    max_tokens=1024,\n",
    "    stream=False,  # Non-streaming\n",
    "    tools=tools,\n",
    ")\n",
    "print(\"Non-stream response:\")\n",
    "print(response_non_stream)\n",
    "print(\"==== content ====\")\n",
    "print(response_non_stream.choices[0].message.content)\n",
    "print(\"==== tool_calls ====\")\n",
    "print(response_non_stream.choices[0].message.tool_calls)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Handle Tools\n",
    "When the engine determines it should call a particular tool, it will return arguments or partial arguments through the response. You can parse these arguments and later invoke the tool accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "name_non_stream = response_non_stream.choices[0].message.tool_calls[0].function.name\n",
    "arguments_non_stream = (\n",
    "    response_non_stream.choices[0].message.tool_calls[0].function.arguments\n",
    ")\n",
    "\n",
    "print(f\"Final streamed function call name: {name_non_stream}\")\n",
    "print(f\"Final streamed function call arguments: {arguments_non_stream}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Streaming Request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Streaming mode test\n",
    "print(\"Streaming response:\")\n",
    "response_stream = client.chat.completions.create(\n",
    "    model=model_name,\n",
    "    messages=messages,\n",
    "    temperature=0,\n",
    "    top_p=0.95,\n",
    "    max_tokens=1024,\n",
    "    stream=True,  # Enable streaming\n",
    "    tools=tools,\n",
    ")\n",
    "\n",
    "texts = \"\"\n",
    "tool_calls = []\n",
    "name = \"\"\n",
    "arguments = \"\"\n",
    "for chunk in response_stream:\n",
    "    if chunk.choices[0].delta.content:\n",
    "        texts += chunk.choices[0].delta.content\n",
    "    if chunk.choices[0].delta.tool_calls:\n",
    "        tool_calls.append(chunk.choices[0].delta.tool_calls[0])\n",
    "print(\"==== Text ====\")\n",
    "print(texts)\n",
    "\n",
    "print(\"==== Tool Call ====\")\n",
    "for tool_call in tool_calls:\n",
    "    print(tool_call)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Handle Tools\n",
    "When the engine determines it should call a particular tool, it will return arguments or partial arguments through the response. You can parse these arguments and later invoke the tool accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parse and combine function call arguments\n",
    "arguments = []\n",
    "for tool_call in tool_calls:\n",
    "    if tool_call.function.name:\n",
    "        print(f\"Streamed function call name: {tool_call.function.name}\")\n",
    "\n",
    "    if tool_call.function.arguments:\n",
    "        arguments.append(tool_call.function.arguments)\n",
    "\n",
    "# Combine all fragments into a single JSON string\n",
    "full_arguments = \"\".join(arguments)\n",
    "print(f\"streamed function call arguments: {full_arguments}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a Tool Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is a demonstration, define real function according to your usage.\n",
    "def get_current_weather(city: str, state: str, unit: \"str\"):\n",
    "    return (\n",
    "        f\"The weather in {city}, {state} is 85 degrees {unit}. It is \"\n",
    "        \"partly cloudly, with highs in the 90's.\"\n",
    "    )\n",
    "\n",
    "\n",
    "available_tools = {\"get_current_weather\": get_current_weather}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Execute the Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages.append(response_non_stream.choices[0].message)\n",
    "\n",
    "# Call the corresponding tool function\n",
    "tool_call = messages[-1].tool_calls[0]\n",
    "tool_name = tool_call.function.name\n",
    "tool_to_call = available_tools[tool_name]\n",
    "result = tool_to_call(**(json.loads(tool_call.function.arguments)))\n",
    "print(f\"Function call result: {result}\")\n",
    "# messages.append({\"role\": \"tool\", \"content\": result, \"name\": tool_name})\n",
    "messages.append(\n",
    "    {\n",
    "        \"role\": \"tool\",\n",
    "        \"tool_call_id\": tool_call.id,\n",
    "        \"content\": str(result),\n",
    "        \"name\": tool_name,\n",
    "    }\n",
    ")\n",
    "\n",
    "print(f\"Updated message history: {messages}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Send Results Back to Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_response = client.chat.completions.create(\n",
    "    model=model_name,\n",
    "    messages=messages,\n",
    "    temperature=0,\n",
    "    top_p=0.95,\n",
    "    stream=False,\n",
    "    tools=tools,\n",
    ")\n",
    "print(\"Non-stream response:\")\n",
    "print(final_response)\n",
    "\n",
    "print(\"==== Text ====\")\n",
    "print(final_response.choices[0].message.content)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
